// 页表以及虚拟内存管理
#include "common.h"
#include "memlayout.h"
#include "riscv.h"
#include "spinlock.h"
#include "proc.h"

// 内核页表(该内核内核页表采用直接映射的方式)
pagetable_t kernel_pagetable;

extern char etext[];      // 内核结束位置
extern char trampoline[];    // trampoline开始位置

// 逐级查询页表，返回页表项
// pgtbl: 页表; va: 虚拟地址; alloc: 是否分配页表
pte_t* walk(pagetable_t pagetable, uint64 va, int alloc) {
    if(va >= MAXVA) {
        panic("超出最大虚拟地址");
    }
    // 逐级查询页表
    for(int level = 2; level > 0; level--) {
        pte_t *pte = &pagetable[VPN(va, level)]; // 根据页表索引在页表中查找页表项
        if(*pte & PTE_V) { // 如果页表项有效，则从中提取下一级页表的物理地址
            pagetable = (pagetable_t)PTE2PA(*pte);
        } else {
            // 如果页表项无效且不需要分配页表项，或物理页面分配失败，则返回0
            if(!alloc || (pagetable = (pte_t*)kalloc()) == 0) {
                return 0;
            }
            memset(pagetable, 0, PGSIZE); // 清零页表
            // 否则将分配的物理页转换为页表项，并置为有效
            *pte = PA2PTE(pagetable) | PTE_V;
        }
    }
    // 返回最后一级页表的页表项
    return &pagetable[VPN(va, 0)];
}
// 在页表中添加虚拟地址到物理地址的映射
// va: 虚拟地址; pa: 物理地址; size: 映射大小; perm: 权限
// 返回值：0表示成功，-1表示失败
int mappage(pagetable_t kpgtbl, uint64 va, uint64 pa, uint64 size, int perm) {
    if(va % PGSIZE != 0) {
        panic("映射失败，虚拟地址未对齐！");
    }
    if(size % PGSIZE != 0) {
        panic("映射失败，映射大小未对齐！");
    }
    if(size == 0) {
        panic("映射失败，映射大小为0！");
    }
    uint64 curr, last;
    pte_t *pte;
    curr = va;                   // 当前映射的页面
    last = va + size - PGSIZE;   // 最后一个映射的页面
    for(;;) {
        // 分配页表失败，则返回-1
        if((pte = walk(kpgtbl, curr, 1)) == 0) {
            return -1;
        }
        // 如果返回了已经映射的页表项，则报错
        if(*pte & PTE_V) {
            panic("映射失败，该页面已经存在映射！");
        }
        // 否则建立页表项与页表的映射
        *pte = PA2PTE(pa) | perm | PTE_V;
        if(curr == last) {
            break;
        }
        curr += PGSIZE;
        pa += PGSIZE;
    }
    return 0;
}

// 在内核页表中添加映射
void kvmmap(pagetable_t kpgtbl, uint64 va, uint64 pa, uint64 size, int perm) {
    if(mappage(kpgtbl, va, pa, size, perm) != 0) {
        panic("内核页表映射失败！");
    }
}
// 为内核创建一个直接映射的页表
pagetable_t kvmmake(void) {
    pagetable_t kpgtbl;     // 内核页表的地址

    kpgtbl = (pagetable_t)kalloc();        // 分配一页物理内存作为根页表
    memset(kpgtbl, 0, PGSIZE);             // 清零页表
    // 直接映射uart，权限为读写
    kvmmap(kpgtbl, UART0, UART0, PGSIZE, PTE_R | PTE_W);
    // 直接映射Virtio mmio disk接口，权限为读写
    kvmmap(kpgtbl, VIRTIO0, VIRTIO0, PGSIZE, PTE_R | PTE_W);
    // 直接映射plic，大小为128MB（即128K个页面）权限为读写
    kvmmap(kpgtbl, PLIC, PLIC, 0x4000000, PTE_R | PTE_W);
    // 直接映射内核代码，权限为可读可执行
    kvmmap(kpgtbl, KERNBASE, KERNBASE, (uint64)etext - KERNBASE, PTE_R | PTE_X);
    // 直接映射本内核数据和可用的物理内存，权限为读写
    kvmmap(kpgtbl, (uint64)etext, (uint64)etext, PHYSTOP - (uint64)etext, PTE_R | PTE_W);
    // 直接映射trampoline，虚拟地址为最高地址，物理地址为内核代码结束处，权限为可读可执行
    kvmmap(kpgtbl, TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);

    // 映射每个进程的内核栈
    // 注意：如果忽略将导致在切换到新进程的上下文后由于sp指针没有正确映射而陷入异常
    procstacksmap(kpgtbl);
    
    return kpgtbl;
}

// 初始化内核页表
void kvminit(void) {
    kernel_pagetable = kvmmake();
    printf("内核页表初始化完成！\n");
}

// 切换到内核页表并启动分页
void kvminithart() {
    // 第一次刷新TLB，确保所有写都刷新到内存
    sfence_vma();
    uint64 x = SATP_MODE_SV39;      // 设置为Sv39模式
    x |= ((uint64)kernel_pagetable >> 12);  // 设置根页表地址
    w_csr(SATP, x);                 // 写入satp寄存器
    // 第二次刷新TLB，确保TLB中的缓存被清空
    sfence_vma();
    printf("hart%d：切换到内核页表完成！\n", r_tp());
}

// 创建一个空进程页表
// 返回值：进程页表物理地址，0表示失败
pagetable_t uvmmake() {
    pagetable_t upgtbl;
    // 为该页表分配物理内存并清空页表项
    upgtbl = (pagetable_t)kalloc();
    if(upgtbl == 0) {
        return 0;
    }
    memset(upgtbl, 0, PGSIZE);
    return upgtbl;
}

// 解除页表从va开始的npages个虚拟地址的映射并选择性释放物理页
// va: 虚拟地址; npages: 解除映射的页面数; dofree: 是否释放物理页
void uvmunmap(pagetable_t upgtbl, uint64 va, uint64 npages, int isfree) {
    if((va % PGSIZE) != 0) {
        panic("解除映射失败，虚拟地址未对齐！");
    }
    uint64 curr;
    pte_t *pte;
    for(curr = va; curr < va + npages * PGSIZE; curr += PGSIZE) {
        // 逐级查找该虚拟地址对应的页表项，但不创建
        if((pte = walk(upgtbl, curr, 0)) == 0) {
            panic("解除映射失败，没有映射！");
        }
        // 如果页表项只设置了有效位权限，说明该页是中间页
        // 而该函数的作用主要是解除映射，即解除最后一级页表映射
        if(PTE_FLAGS(*pte) == PTE_V) {
            panic("解除映射失败，试图释放中间页！");
        }
        // 释放物理页
        if(isfree) {
            uint64 pa = PTE2PA(*pte);
            kfree((void*)pa);
        }
        // 清空即解除映射
        pte = 0;
    }
}

// 递归释放页表
// 只释放中间页表，而不释放叶子页，即实际占用的内存
// 即使释放也是由uvmunmap负责
void freewalk(pagetable_t pagetable) {
    for(int i = 0; i < 512; i++) {
        pte_t pte = pagetable[i];
        // 如果该页表项有效且是中间页，则递归释放子页表
        if((pte & PTE_V) && (pte & (PTE_R | PTE_W | PTE_X)) == 0) {
            uint64 pa = PTE2PA(pte);
            freewalk((pagetable_t)pa);
            pagetable[i] = 0;
        } else if(pte & PTE_V) {
            panic("释放页表失败，试图释放叶子页！");
        }
    }
    // 释放整个页表
    kfree((void*)pagetable);
}

// 释放一个进程页表，遍历释放每一个页表和物理页
// upgtbl: 进程页表; sz: 进程使用的物理空间大小
void uvmfree(pagetable_t upgtbl, uint64 sz) {
    // 如果进程使用了物理内存，则解除地址映射并释放物理页
    if(sz > 0) {
        uvmunmap(upgtbl, 0, PGROUNDUP(sz) / PGSIZE, 1);
    }
    // 递归释放页表
    freewalk(upgtbl);
}

// 为第一个用户进程分配第一个页面（虚拟地址为0）并写入指令
void uvmfirst(pagetable_t upgtbl, uint8 *src, uint32 sz) {
    if(sz > PGSIZE) {
        panic("第一个进程页面分配失败，大小超过一页！");
    }
    char *mem = kalloc();
    memset(mem, 0, PGSIZE);
    mappage(upgtbl, 0, (uint64)mem, PGSIZE, PTE_R | PTE_W | PTE_X | PTE_U);
    memmove(mem, src, sz);
}